import os
import uuid
import json
import logging
from typing import List, Dict, Optional
import re

from label_studio_ml.model import LabelStudioMLBase
from label_studio_ml.response import ModelResponse
from openai import OpenAI
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()
logger = logging.getLogger(__name__)
# 验证env是否生效
# logger.info(f"读取的 LS 地址是：{os.environ.get('LABEL_STUDIO_URL')}")

# 初始化 OpenAI 客户端
client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),  # 请替换为你的API密钥
    base_url="https://api.deepseek.com",  # deepseek 作为兼容 OpenAI API 的替代服务
)


class NewModel(LabelStudioMLBase):
    """定义后端模型"""

    # 获取标签列表
    def get_labels(self):
        li = self.label_interface
        from_name, _, _ = li.get_first_tag_occurence('Labels', 'Text')
        tag = li.get_tag(from_name)
        return tag.labels

    def setup(self):
        """配置模型参数"""
        self.set("model_version", "0.0.1")

    def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> ModelResponse:
        """执行预测逻辑"""
        li = self.label_interface
        from_name, to_name, value = li.get_first_tag_occurence('Labels', 'Text')
        labels = self.get_labels()
        predictions = []

        for task in tasks:
            data_value = task['data'].get(value)
            if not data_value:
                logger.warning(f"任务 {task['id']} 缺少字段 '{value}'")
                continue

            try:
                if isinstance(data_value, str) and data_value.startswith('/data/'):
                    # 如果是文件路径，读取文件内容
                    file_url = task['data'][value]
                    try:
                        local_file = self.get_local_path(file_url, task_id=task['id'])
                        print("本地文件路径:", local_file)
                        # 读取文本内容
                        with open(local_file, 'r', encoding='utf-8') as f:
                            text = f.read()

                    except Exception as e:
                        logger.error(f"读取任务文件失败：{file_url}, 错误信息: {e}")
                        continue
                else:
                    # 如果是文本字符串，提取文本
                    text = self.preload_task_data(task, task['data'][value])
            except Exception as e:
                logger.error(f"任务 {task['id']} 的数据处理失败：{e}")
                continue

            # 构造 system prompt（让 DeepSeek 按预期返回 JSON 格式）
            system_prompt = f"""你是一名命名实体识别 (NER) 助手。请根据输入文本进行命名实体识别（NER），并从以下标签中选择最合适的标签：{labels}。
                    请返回一个 JSON 格式，每个花括号中代表一个命名实体，格式如下：
                        [{{
                            "start": <实体在文本中的起始位置>,
                            "end": <结束位置>,
                            "text": "<实体名称>",
                            "labels": ["<标签>"],
                            "score": <置信度分数>
                        }},
                        ...
                        ]
                    不要包含任何解释或自然语言，只输出这个 JSON 数组。"""

            try:
                # 调用 DeepSeek 进行推理
                response = client.chat.completions.create(
                    model="deepseek-chat",  # 可以使用模型名称，例如 "deepseek-chat"
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": text}
                    ],
                    response_format={
                        'type': 'json_object'
                    },
                    temperature=0
                )

                # 获取并解析返回内容
                result_text = response.choices[0].message.content

                try:
                    result = json.loads(result_text)
                except json.JSONDecodeError as e:
                    logger.error(f"DeepSeek 返回的内容无法解析为 JSON: {e}")
                    continue

                # 如果模型返回了 {"entities": [...]} 格式，取出列表
                if isinstance(result, dict) and "entities" in result:
                    result = result["entities"]

                # 生成多个实体的预测结果
                result_entities = []
                used_offsets = {}
                for entity in result:

                    # 获取实体信息
                    entity_text = entity.get("text")
                    entity_labels = entity.get("labels")
                    entity_score = entity.get("score", 1.0)

                    # fallback 的 start/end
                    entity_start = entity.get("start", -1)
                    entity_end = entity.get("end", -1)

                    # 使用 finditer 查找所有匹配位置
                    matches = list(re.finditer(re.escape(entity_text), text))

                    if matches:
                        # 初始化使用记录
                        if entity_text not in used_offsets:
                            used_offsets[entity_text] = 0
                        match_index = used_offsets[entity_text]
                        if match_index < len(matches):
                            match = matches[match_index]
                            entity_start = match.start()
                            entity_end = match.end()
                            used_offsets[entity_text] += 1
                        else:
                            logger.warning(
                                f"实体 '{entity_text}' 超出匹配次数，仅使用原始偏移 start={entity_start}, end={entity_end}")
                    else:
                        logger.warning(f"未找到实体 '{entity_text}' 的任何匹配位置，使用原始偏移")

                    if isinstance(entity_labels, str):
                        entity_labels = [entity_labels]

                    # 生成唯一的 ID
                    # entity_id = str(uuid.uuid4())

                    # 构建实体结果
                    result_entities.append({

                        "from_name": from_name,
                        "to_name": to_name,
                        "type": "labels",
                        "value": {
                            "start": entity_start,
                            "end": entity_end,
                            "text": entity_text,
                            "score": entity_score,
                            "labels": entity_labels  # 取第一个标签
                        }
                    })

                # 封装为预测结果
                prediction_result = {
                    "model_version": self.get("model_version"),
                    "result": result_entities
                }

                predictions.append(prediction_result)

            except Exception as e:
                logger.error(f"DeepSeek 推理失败：{e}")
                continue

        return ModelResponse(predictions=predictions)

    def fit(self, event, data, **kwargs):
        """
        This method is called each time an annotation is created or updated.
        """
        old_data = self.get('my_data')
        old_model_version = self.get('model_version')
        print(f'Old data: {old_data}')
        print(f'Old model version: {old_model_version}')

        # store new data to the cache
        self.set('my_data', 'my_new_data_value')
        self.set('model_version', 'my_new_model_version')
        print(f'New data: {self.get("my_data")}')
        print(f'New model version: {self.get("model_version")}')

        print('fit() completed successfully.')
